import torch
from data import prep_data
from models import Net
# 导入必要的库
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import train_test_split
import warnings
import os
from sklearn.preprocessing import OneHotEncoder, LabelEncoder

# 导入数据
# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))


# 使用绝对路径
model_path = os.path.join(current_dir, "trained_model.pth")

def predict_mental_health(sample_data, model_path=model_path):
    # 加载模型
    input_size = sample_data.shape[1]  # 获取输入特征的维度
    model = Net(input_size)  # 创建神经网络模型
    model.load_state_dict(torch.load(model_path))  # 请确保训练好的模型文件 'trained_model.pth' 位于相同目录下
    # 设置模型为评估模式
    model.eval()

    # 进行预测
    with torch.no_grad():
        output = model(sample_data)
        prediction = (output > 0.5).item()  # 将概率值大于0.5的转换为1，否则为0，并获取预测结果

    return prediction